import random

import av
import torch
from cv2 import cv2

from common.log.logUtils import logUtils
from models.experimental import attempt_load
from rtcCVTransform.videoTrans_Base import videoTransformBase
import numpy as np

from utils.datasets import letterbox
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box
from utils.torch_utils import select_device
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


class videoTransYolo(videoTransformBase):

    def __init__(self, weights, device):
        # assign device
        self.device = select_device(device)
        # assign precision
        self.half = self.device.type != 'cpu'  # half precision only supported on CUDA
        self.augment = True
        # 加载检测模型
        self.model = attempt_load(weights, map_location=self.device)

        if self.half:
            self.model.half()  # to FP16

        self.img_size=640
        self.stride = int(self.model.stride.max())

        # Get names and colors
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]

        logUtils.info("load yolo detect model finished....")
        logUtils.info("names:"+str(self.names))

    def transform(self, frame: av.VideoFrame) -> np.ndarray:
        img = frame.to_ndarray(format="bgr24")
        img=self.getYoloDetectResult(img)
        return img

    def getYoloDetectResult(self, img0: np.ndarray):

        img = letterbox(img0, self.img_size, stride=self.stride)[0]

        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)
        # 1. convert numpy to tensor
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # 2.predict with model
        pred = self.model(img, augment=self.augment)[0]

        # Apply NMS
        conf_thres=0.4
        iou_thres=0.45
        classes=None
        agnostic_nms=False
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)

        for i, det in enumerate(pred):  # detections per image

            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

                # Write results
                for *xyxy, conf, cls in reversed(det):
                        c = int(cls)  # integer class
                        label = f'{self.names[c]} {conf:.2f}'
                        plot_one_box(xyxy, img0, label=label, color=self.colors[c], line_thickness=3)
                # cv2.imshow("draw box",img0)
                # cv2.waitKey(0)

        return img0


# yoloModel = videoTransYolo("../data/models/yolov5s.pt", 'cpu')
# #
# img=cv2.imread("../data/img/share.jpg")
# yoloModel.getYoloDetectResult(img)
